import openai
from sglang.utils import print_highlight
import readline  # 解决命令行编辑中文的问题

port = 31266
client = openai.Client(base_url=f"http://127.0.0.1:{port}/v1", api_key="None")

# 初始化对话历史
messages = [{"role": "system", "content": "用户要开始和你聊天了，准备好哦"}]

print("开始对话 (输入 'exit' 退出):")
while True:
    # 获取用户输入
    user_input = input("\n[我]: ")
    
    # 退出条件
    if user_input.lower() in ["exit", "quit"]:
        print("退出对话")
        out_path = "chat_history.txt"
        with open(out_path, 'w') as file:
            for message in messages:
                file.write(f'[{message["role"]}]: {message["content"]}\n')
        print(f"聊天记录保存至: {out_path}")
        break
    
    # 添加用户消息到历史
    messages.append({"role": "user", "content": user_input})
    
    try:
        # 发送请求
        response = client.chat.completions.create(
            model="../../models/Qwen3-0.6B",
            messages=messages,
            temperature=0,
            max_tokens=1024,
        )
        
        # 获取模型回复
        assistant_reply = response.choices[0].message.content
        print(f"\n[助手]: {assistant_reply}")
        
        # 添加助手回复到历史
        messages.append({"role": "assistant", "content": assistant_reply})
    
    except Exception as e:
        print(f"退出: {e}")
        for message in messages:
            print(message)
        # 移除最后一条用户消息（允许重新输入）
        messages.pop()